from datetime import date, datetime
import os
import shutil
import jieba
from src.database.dao_factory import DAOFactory
from src.utils.util import Util
from whoosh.index import create_in, exists_in, open_dir
from whoosh.fields import Schema, ID, TEXT, DATETIME, KEYWORD, NUMERIC
from whoosh.qparser import QueryParser, MultifieldParser
from whoosh.analysis import Tokenizer, Token, NgramAnalyzer
from whoosh.query import DateRange

from src.database.dao.video_dao import VideoDAO
from src.database.models import Video
from src.utils.logger import log_info
indexdir = os.path.join(os.path.expanduser("~"), "VideoClassifier", "whoosh_index")

class WhooshProcessor:
    def __init__(self, indexdir=indexdir):
        self.indexdir = indexdir
        self.dao_factory = DAOFactory.get_instance()
        self.schema = Schema(
            video_id=ID(stored=True, unique=True),
            description=TEXT(stored=True, analyzer=JiebaAnalyzer()),
            category=TEXT(stored=True),
            downloaded_date=DATETIME(stored=True),
            tags=KEYWORD(stored=True, commas=True, lowercase=True, analyzer=NgramAnalyzer(1, 3)),
            duration=NUMERIC(stored=True, sortable=True)  # Add duration field
        )

        if exists_in(indexdir):
            log_info(f"Loading Whoosh index from {indexdir}")
            self.ix = open_dir(indexdir)
        else:
            log_info(f"Creating new Whoosh index at {indexdir}")
            os.makedirs(indexdir, exist_ok=True)
            self.ix = create_in(indexdir, self.schema)        

    def add_video(self, videos: list[Video]):
        writer = self.ix.writer()
        added_count = 0
        skipped_count = 0
        
        for video in videos:
            # Skip videos without .mp4 extension
            if not video._local_path or not video._local_path.lower().endswith('.mp4')\
                or '\\_固定水印或字幕\\' in video._local_path:
                skipped_count += 1
                continue
                
            # Extract description and tags from title
            description = video.title
            tags = []
            
            # Parse hashtags from title
            if '#' in description:
                parts = description.split('#')
                # First part is the description
                description = parts[0].strip()
                # Remaining parts are hashtags
                for part in parts[1:]:
                    if part.strip():
                        tags.append(part.strip().lower())
            
            writer.add_document(
                video_id=str(video.id), 
                description=description,
                category=video.category, 
                downloaded_date=video.download_at,
                tags=",".join(tags),  # Join tags with commas for KEYWORD field
                duration=video.duration
            )
            log_info(f"DURATION {video.duration}")
            added_count += 1
            
        writer.commit()
        
        # Log summary information
        if added_count > 0 or skipped_count > 0:
            log_info(f"索引添加: {added_count} 个MP4视频文件, 跳过 {skipped_count} 个非MP4文件")

    def search(self, category, query_str, start_date: str=None, end_date: str=None,
               max_duration=None, limit=3):
            
        # 先用jieba分词
        limit = int(limit)
        log_info("Start_date:", start_date)
        log_info("End_date:", end_date)
        expanded_words = Util.expand_keywords(query_str)
        # 基于分词结果做扩展，得到更多相关词
        if expanded_words:
            query_str = " OR ".join(expanded_words)
        else:
            query_str = " OR ".join(query_str)  # 兜底直接用分词
        
        results = []
        with self.ix.searcher() as searcher:
            # Create a query that combines category filter and description search
            category_query = QueryParser("category", self.ix.schema).parse(f'"{category}"')
            if not query_str:
                query = category_query
            else:
                _parser = MultifieldParser(["description", "tags"], schema=self.ix.schema)
                description_query = _parser.parse(query_str)
                query = category_query & description_query
            # Add date filtering if specified
            if start_date or end_date:
                start_dt = datetime.strptime(start_date, "%Y-%m-%d %H:%M:%S") if start_date else None
                end_dt = datetime.strptime(end_date, "%Y-%m-%d %H:%M:%S") if end_date else None
                date_query = DateRange("downloaded_date", start_dt, end_dt)
                query = query & date_query
                        
            if max_duration is not None:
                # Add duration filtering if specified
                from whoosh.query import NumericRange
                duration_query = NumericRange("duration", 0, max_duration)
                query = query & duration_query
            log_info(f"Final search query: {query}")
            for hit in searcher.search(query, limit=limit):
                video_id = hit['video_id']
                description = hit['description']
                results.append({
                    "video_id": video_id,
                    "description": description,
                    "score": hit.score,
                    "category": hit['category'],
                })
        return results
        
    def fetch_all(self):
        with self.ix.searcher() as searcher:
            results = searcher.documents()
            for hit in results:
                log_info(hit['video_id'], hit['description'])

    def init_videos(self):
        # Reset index folder
        if os.path.exists(self.indexdir):
            shutil.rmtree(self.indexdir)
        os.makedirs(self.indexdir, exist_ok=True)
        self.ix = create_in(self.indexdir, self.schema)
        
        # Load videos
        self.dao_factory.get_video_dao().process_in_batches(
            callback=self.add_video,
            batch_size=1000
        )

    def _search_tags(self, category, query_str=None, limit=10):
        """
        Search for tags based on input query, filtered by category.
        
        :param category: Category to filter by
        :param query_str: Query string to match against tags
        :param limit: Maximum number of tags to return
        :return: List of matching tags
        """
        results = set()  # Use a set to avoid duplicates
        try:
            with self.ix.searcher() as searcher:
                # Create category query
                category_query = QueryParser("category", self.ix.schema).parse(f'"{category}"')
                
                if query_str and query_str.strip():
                    # Search for tags that contain the query string, filtered by category
                    tag_query = QueryParser("tags", self.ix.schema).parse(query_str.lower())
                    combined_query = category_query & tag_query
                    
                    hits = searcher.search(combined_query, limit=limit*3)
                    
                    # Extract and sort tags by relevance
                    tag_scores = {}
                    for hit in hits:
                        if 'tags' in hit and hit['tags']:
                            for tag in hit['tags'].split(','):
                                if tag and query_str.lower() in tag.lower():
                                    if tag in tag_scores:
                                        tag_scores[tag] += hit.score
                                    else:
                                        tag_scores[tag] = hit.score
                    
                    # Sort tags by relevance score
                    sorted_tags = sorted(tag_scores.items(), key=lambda x: x[1], reverse=True)
                    results = [tag for tag, _ in sorted_tags[:limit]]
                else:
                    # Get most recent tags from recent videos in the specified category
                    hits = searcher.search(category_query, sortedby="downloaded_date", reverse=True, limit=limit*5)
                    
                    # Extract tags
                    for hit in hits:
                        if 'tags' in hit and hit['tags']:
                            for tag in hit['tags'].split(','):
                                if tag and len(results) < limit:
                                    results.add(tag)
                    
                    # Convert set to list
                    results = list(results)[:limit]
        except Exception as e:
            log_info(f"Error searching tags: {str(e)}")
        
        return list(results)

    def _get_popular_tags(self, category, limit=10):
        """
        Get most frequently used tags across videos in the specified category.
        
        :param category: Category to filter by
        :param limit: Maximum number of tags to return
        :return: List of most popular tags in the category
        """
        tag_counts = {}
        try:
            with self.ix.searcher() as searcher:
                # Create category query
                category_query = QueryParser("category", self.ix.schema).parse(f'"{category}"')
                
                # Search for all documents in the category
                hits = searcher.search(category_query, limit=None)
                
                # Count tags in the filtered documents
                for hit in hits:
                    if 'tags' in hit and hit['tags']:
                        for tag in hit['tags'].split(','):
                            if tag:
                                tag = tag.strip()
                                tag_counts[tag] = tag_counts.get(tag, 0) + 1
                
                # Sort tags by frequency
                sorted_tags = sorted(tag_counts.items(), key=lambda x: x[1], reverse=True)
                return [tag for tag, _ in sorted_tags[:limit]]
        except Exception as e:
            log_info(f"Error getting popular tags: {str(e)}")
            return []
        
    def get_tag_suggestions(self, category, query_str=None, limit=10):
        """
        Main method for getting tag suggestions.
        Returns relevant tags if query is provided, or popular tags if not.
        
        :param query_str: Optional search query
        :param limit: Maximum number of tags to return
        :return: List of suggested tags
        """
        if query_str and query_str.strip():
            return self._search_tags(category, query_str, limit)
        else:
            return self._get_popular_tags(category, limit)

class JiebaTokenizer(Tokenizer):
    def __call__(self, value, **kwargs):
        t = Token()
        for word in jieba.cut(value, cut_all=False):
            t.original = t.text = word
            t.boost = 1.0
            t.startchar = 0
            t.endchar = len(value)
            t.pos = 0
            yield t

def JiebaAnalyzer():
    return JiebaTokenizer()


